import requests
import difflib
import smtplib
import sqlite3
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from apscheduler.schedulers.blocking import BlockingScheduler
import os
from datetime import datetime
import yaml
import logging
from jinja2 import Environment, FileSystemLoader
from concurrent.futures import ThreadPoolExecutor
import time  # 添加导入time模块

# 配置日志
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

# 读取配置文件
config_yml_path = 'config.yml'
if not os.path.exists(config_yml_path):
    logging.error(f"文件 {config_yml_path} 未找到，请确保文件在同级目录中。")
    raise FileNotFoundError(f"文件 {config_yml_path} 未找到，请确保文件在同级目录中。")

with open(config_yml_path, 'r', encoding='utf-8') as file:
    config = yaml.safe_load(file)

CHECK_INTERVAL = config['CHECK_INTERVAL']  # 每12小时检查一次（单位：秒）
EMAIL_FROM = config['EMAIL_FROM']  # 你的QQ邮箱
EMAIL_TO = config['EMAIL_TO']  # 接收通知的邮箱
SMTP_SERVER = config['SMTP_SERVER']
SMTP_PORT = config['SMTP_PORT']  # QQ邮箱的SSL端口
SMTP_PASSWORD = config['SMTP_PASSWORD']  # QQ邮箱的授权码
DATABASE_FILE = config['DATABASE_FILE']  # SQLite数据库文件
SERVER_HOST = config['SERVER_HOST']
HTML_SERVER_HOST = config['HTML_SERVER_HOST']  # 用于邮件中的链接
SERVER_PORT = config['SERVER_PORT']  # HTTP服务器端口
OUTPUT_DIR = config['OUTPUT_DIR']  # 输出目录
DOMAINS_YML_PATH = config['domains_yml_path']

# 读取 YAML 文件中的目标域名
if not os.path.exists(DOMAINS_YML_PATH):
    logging.error(f"文件 {DOMAINS_YML_PATH} 未找到，请确保文件在同级目录中。")
    raise FileNotFoundError(f"文件 {DOMAINS_YML_PATH} 未找到，请确保文件在同级目录中。")

with open(DOMAINS_YML_PATH, 'r', encoding='utf-8') as file:
    config = yaml.safe_load(file)
TARGET_DOMAINS = config['domains']

# 确保输出目录存在
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

# 初始化数据库
def init_db():
    """初始化数据库"""
    conn = sqlite3.connect(DATABASE_FILE)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS domains (
            domain TEXT PRIMARY KEY,
            subdomains TEXT
        )
    """)
    conn.commit()
    conn.close()

# 获取当前域名列表
def fetch_domains(domain):
    """获取指定域名的子域名列表"""
    domains_set = set()

    # 证书透明度日志（CT Logs）
    try:
        response = requests.get(f"https://crt.sh/?q={domain}&output=json")
        response.raise_for_status()  # 检查HTTP请求是否成功
        ct_logs = response.json()
        for entry in ct_logs:
            name_value = entry['name_value']
            # 处理多行域名的情况
            for name in name_value.split('\n'):
                domains_set.add(name.strip())
        logging.debug(f"CT Logs查询结果: {domains_set}")
    except requests.exceptions.HTTPError as e:
        logging.error(f"CT Logs查询失败: HTTP状态码 {response.status_code}, 错误信息: {e}")
    except requests.exceptions.RequestException as e:
        logging.error(f"CT Logs查询请求失败: {e}")
    except ValueError as e:
        logging.error(f"CT Logs查询返回的数据解析失败: {e}")
    except Exception as e:
        logging.error(f"CT Logs查询未知错误: {e}")

    return list(domains_set)

# 对比域名列表
def compare_domains(old_domains, new_domains):
    """对比两个域名列表，返回新增和删除的域名"""
    old_set = set(old_domains)
    new_set = set(new_domains)
    added_domains = new_set - old_set
    removed_domains = old_set - new_set
    changes = []
    for domain in added_domains:
        changes.append(f"+ {domain}")
    for domain in removed_domains:
        changes.append(f"- {domain}")
    return changes

# 发送邮件通知
def send_email(changes, domain, initial=False):
    """发送邮件通知"""
    subject = f"域名 {domain} 变更通知"
    if initial:
        body_text = "首次监控结果:\n" + "\n".join(changes[:10])
        if len(changes) > 10:
            body_text += f"\n... 还有 {len(changes) - 10} 条变化"
    else:
        if not changes:
            body_text = "近期无新增域名，持续监控中"
        else:
            # 只发送新增的域名
            added_changes = [change for change in changes if change.startswith('+')]
            body_text = "\n".join(added_changes[:10])  # 只发送前10条变化
            if len(added_changes) > 10:
                body_text += f"\n... 还有 {len(added_changes) - 10} 条变化"

    # 生成HTML内容
    html_content = generate_html(changes, domain, initial)
    file_name = save_html(html_content, domain)
    link = f"http://{HTML_SERVER_HOST}:{SERVER_PORT}/{file_name}"
    body_text += f"\n\n更多详细信息请访问: <a href='{link}'>{link}</a>"

    # 创建邮件对象
    message = MIMEMultipart()
    message['From'] = EMAIL_FROM
    message['To'] = EMAIL_TO
    message['Subject'] = subject

    html_part = MIMEText(body_text, "html", "utf-8")
    message.attach(html_part)

    try:
        # 使用SSL连接SMTP服务器
        with smtplib.SMTP_SSL(SMTP_SERVER, SMTP_PORT) as server:
            server.login(EMAIL_FROM, SMTP_PASSWORD)
            server.sendmail(EMAIL_FROM, [EMAIL_TO], message.as_string())
        logging.info("邮件发送成功！")
    except smtplib.SMTPAuthenticationError:
        logging.error("SMTP身份验证失败，请检查邮箱地址和授权码")
        logging.error("请确保您已启用SMTP服务并获取正确的授权码。更多信息请参考：https://service.mail.qq.com/detail/0/53")
    except smtplib.SMTPException as e:
        logging.error(f"SMTP异常: {e}")
    except Exception as e:
        logging.error(f"邮件发送失败: {e}")

# 生成HTML内容
def generate_html(changes, domain, initial=False):
    """生成HTML内容"""
    env = Environment(loader=FileSystemLoader('.'))
    template = env.get_template('email_template.html')
    if initial:
        # 获取所有域名
        all_domains = [change.split(' ', 1)[1] for change in changes if change.startswith('+')]
        change_count = len(all_domains)  # 修改这里，使用所有域名的数量
        message = f"首次监控结果: 检测到 {change_count} 个域名:"
    else:
        change_count = sum(line.startswith('+') for line in changes)
        # 获取新增域名
        added_subdomains = [change.split(' ', 1)[1] for change in changes if change.startswith('+')]
        all_domains = added_subdomains
        if change_count == 0:
            message = "近期无新增域名，持续监控中"
        else:
            message = f"检测到 {change_count} 个域名变更:"
    domain_dict = {domain: {'count': len(all_domains), 'subdomains': all_domains}}
    return template.render(domain=domain, message=message, domains=domain_dict)

def generate_combined_html(all_changes, initial=False):
    """生成合并后的HTML内容"""
    env = Environment(loader=FileSystemLoader('.'))
    template = env.get_template('email_template.html')
    all_domains = {}
    for domain, changes in all_changes.items():
        added_subdomains = [change.split(' ', 1)[1] for change in changes if change.startswith('+')]
        if initial:
            all_domains[domain] = {'count': len(added_subdomains), 'subdomains': added_subdomains}
        else:
            if added_subdomains:
                all_domains[domain] = {'count': len(added_subdomains), 'subdomains': added_subdomains}
    change_count = sum(len(changes) for changes in all_changes.values())
    if initial:
        message = f"首次监控结果: 检测到 {change_count} 个域名:"
    else:
        if change_count == 0:
            message = "近期无新增域名，持续监控中"
        else:
            message = f"检测到 {change_count} 个域名变更:"
    return template.render(domain="Combined", message=message, domains=all_domains)

# 保存HTML文件
def save_html(content, domain):
    """保存HTML文件"""
    current_time = datetime.now().strftime("%Y%m%d_%H%M")
    file_name = f"{current_time}-{domain}.html"
    file_path = os.path.join(OUTPUT_DIR, file_name)
    with open(file_path, "w", encoding="utf-8") as file:
        file.write(content)
    logging.info(f"HTML文件已保存为 {file_path}")
    return file_name

# 存储域名数据
def save_domains(domain, subdomains):
    """存储域名数据"""
    unique_subdomains = list(set(subdomains))
    with sqlite3.connect(DATABASE_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("""
            INSERT OR REPLACE INTO domains (domain, subdomains)
            VALUES (?, ?)
        """, (domain, ",".join(unique_subdomains)))
        conn.commit()

# 加载域名数据
def load_domains(domain):
    """加载域名数据"""
    with sqlite3.connect(DATABASE_FILE) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT subdomains FROM domains WHERE domain = ?", (domain,))
        result = cursor.fetchone()
    return result[0].split(",") if result else []

def monitor_single_domain(domain, initial=False):
    """监控单个域名"""
    logging.info(f"正在检查域名: {domain}")
    old_domains = load_domains(domain)
    new_domains = fetch_domains(domain)
    changes = compare_domains(old_domains, new_domains)
    save_domains(domain, new_domains)  # 每次执行都存储一次
    if initial:
        changes = [f"+ {domain}" for domain in new_domains if domain not in old_domains]  # 首次运行时只记录新增域名
    else:
        changes = [change for change in changes if change.startswith('+')]  # 非首次运行时只记录新增域名
    return changes

# 新增: 发送合并后的邮件，包括没有检测到变化的情况
def send_combined_email(all_changes, initial=False):
    """发送合并后的邮件"""
    subject = "域名监控通知"
    body_text = ""
    all_domains = []
    has_changes = False
    for domain, changes in all_changes.items():
        if changes:
            body_text += f"域名 {domain} 变更通知:\n"
            # 只发送新增的域名
            added_changes = [change for change in changes if change.startswith('+')]
            body_text += "\n".join(added_changes[:10])  # 只发送前10条变化
            if len(added_changes) > 10:
                body_text += f"\n... 还有 {len(added_changes) - 10} 条变化\n"
            all_domains.extend([change.split(' ', 1)[1] for change in added_changes if change.startswith('+')])
            has_changes = True
        else:
            body_text += f"域名 {domain} 未检测到新增域名\n"

    all_domains = list(set(all_domains))
    change_count = sum(len(changes) for changes in all_changes.values())
    if initial:
        message = f"首次监控结果: 检测到 {change_count} 个域名:"
    else:
        if not has_changes:
            body_text = "近期无新增域名，持续监控中"
            message = "近期无新增域名，持续监控中"
        else:
            message = f"检测到 {change_count} 个域名变更:"

    # 生成HTML内容
    html_content = generate_combined_html(all_changes, initial)
    file_name = save_html(html_content, "combined")
    link = f"http://{HTML_SERVER_HOST}:{SERVER_PORT}/{file_name}"
    body_text += f"\n\n更多详细信息请访问: <a href='{link}'>{link}</a>"

    # 创建邮件对象
    message = MIMEMultipart()
    message['From'] = EMAIL_FROM
    message['To'] = EMAIL_TO
    message['Subject'] = subject

    html_part = MIMEText(body_text, "html", "utf-8")
    message.attach(html_part)

    try:
        # 使用SSL连接SMTP服务器
        with smtplib.SMTP_SSL(SMTP_SERVER, SMTP_PORT) as server:
            server.login(EMAIL_FROM, SMTP_PASSWORD)
            server.sendmail(EMAIL_FROM, [EMAIL_TO], message.as_string())
        logging.info("邮件发送成功！")
    except smtplib.SMTPAuthenticationError:
        logging.error("SMTP身份验证失败，请检查邮箱地址和授权码")
        logging.error("请确保您已启用SMTP服务并获取正确的授权码。更多信息请参考：https://service.mail.qq.com/detail/0/53")
    except smtplib.SMTPException as e:
        logging.error(f"SMTP异常: {e}")
    except Exception as e:
        logging.error(f"邮件发送失败: {e}")

# 监控所有域名
def monitor_domains():
    """监控所有域名"""
    all_changes = {}
    with ThreadPoolExecutor() as executor:
        results = executor.map(lambda domain: monitor_single_domain(domain, initial=True), TARGET_DOMAINS)  # 首次运行时设置initial=True
        for domain, changes in zip(TARGET_DOMAINS, results):
            all_changes[domain] = changes

    # 首次运行时发送邮件，即使没有变化
    logging.info("首次监控结果，准备发送合并邮件...")
    send_combined_email(all_changes, initial=True)  # 首次运行时设置initial=True

    # 后续运行时设置initial=False
    all_changes = {}
    with ThreadPoolExecutor() as executor:
        results = executor.map(lambda domain: monitor_single_domain(domain, initial=False), TARGET_DOMAINS)
        for domain, changes in zip(TARGET_DOMAINS, results):
            all_changes[domain] = changes

    # 后续运行时发送邮件，即使没有变化
    logging.info("检测结果，准备发送合并邮件...")
    send_combined_email(all_changes, initial=False)  # 后续运行时设置initial=False

# 定时任务
def start_scheduler():
    """启动定时任务"""
    scheduler = BlockingScheduler()
    scheduler.add_job(monitor_domains, 'interval', seconds=CHECK_INTERVAL, max_instances=1)  # 设置最大实例数为1
    scheduler.start()

if __name__ == "__main__":
    init_db()
    # 启动定时任务
    start_scheduler()